

import socket
import struct
import threading

# 定义数据类型
DATA_TYPE_SEND = 0x000A
DATA_TYPE_GET_NEIGHBORS = 0x000B
DATA_TYPE_GET_TOPOLOGY = 0x000C

DATA_TYPE_RECEIVE = 0x00A0
DATA_TYPE_NEIGHBORS = 0x00B0
DATA_TYPE_TOPOLOGY = 0x00C0


class RouterClient:
    def __init__(self, host, port):
        self.host = host
        self.port = port
        self.sock = None
        

    def connect(self):
        """建立TCP连接"""
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.connect((self.host, self.port))
        print(f"Connected to {self.host}:{self.port}")

    def close(self):
        """关闭TCP连接"""
        if self.sock:
            self.sock.close()
            self.sock = None
        print("Connection closed.")

    def create_packet(self, data_type, dst_id, data_length, data):
        """
        创建符合格式的数据包
        数据包格式：
        uint16_t type;    // 类型
        uint16_t subID;   // 固定 0x0002
        uint16_t SRCID;   // 固定 0x0000
        uint16_t DSTID;   // 目的节点 ID
        uint32_t data_len;// 数据长度
        """
        # 打包数据包头部
        header = struct.pack(
            '<HHHHI',  # 格式：2个 uint16_t, 2个 uint16_t, 1个 uint32_t
            data_type,  # 类型 (uint16_t)
            0x0004,     # subID (uint16_t)
            0x0000,     # SRCID (uint16_t)
            dst_id,     # DSTID (uint16_t)
            data_length # 数据长度 (uint32_t)
        )
        # 返回完整的包（头部 + 数据）
        return header + data

    def send_packet(self, data_type, dst_id, data_length=0, data=b''):
        """发送数据包"""
        if not self.sock:
            raise Exception("Connection not established")
        # 创建数据包
        packet = self.create_packet(data_type, dst_id, data_length, data)
        self.sock.sendall(packet)
        print(f"Sent packet: type={hex(data_type)}, dst_id={dst_id}, length={data_length}")

    def receive_packet(self):
        """
        监听一个端口，注意，该端口可能并无连接，当建立连接时及时获取并处理数据
        """
        try:
            # 创建套接字
            print(f"try listening on {self.host}:{self.port}")
            self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.sock.bind((self.host, self.port))
            self.sock.listen(1)  # 开始监听，允许一个连接

            print(f"Listening on {self.host}:{self.port}...")

            # 接受连接
            conn, addr = self.sock.accept()
            print(f"Connected by {addr}")

            # 接收包头（12 字节）
            header = conn.recv(12)
            if not header:
                print("No data received, closing connection.")
                conn.close()
                return None

            # 解析包头
            data_type, sub_id, src_id, dst_id, data_length = struct.unpack('<HHHHI', header)
            print(f"Received header: data_type={data_type}, sub_id={sub_id}, src_id={src_id}, dst_id={dst_id}, data_length={data_length}")

            # 接收数据
            data = conn.recv(data_length)
            if len(data) != data_length:
                raise Exception("Incomplete data received")

            print(f"Received data: {data}")

            # 关闭连接
            conn.close()

            return data_type, dst_id, data_length, data

        except Exception as e:
            print(f"An error occurred: {e}")
            if self.sock:
                self.sock.close()
            return None
